# task_manager.py
import logging
from datetime import datetime
from pathlib import Path
from typing import Dict, Any

from apscheduler.schedulers.background import BackgroundScheduler
from apscheduler.triggers.cron import CronTrigger
from apscheduler.triggers.interval import IntervalTrigger
from selenium.common import TimeoutException, WebDriverException

from image_sender import SyncQyWechatBot
from prometheus_client import start_http_server
from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from PIL import Image
from selenium.webdriver.support.wait import WebDriverWait
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type, wait_fixed
import time
import json
import os


class BrowserManager:
    """浏览器实例全生命周期管理"""

    def __init__(self, browser_config: Dict[str, Any]):
        # 移除默认值，强制要求传入配置
        if not browser_config:
            raise ValueError("Browser config is required")

        self.browser_config = browser_config
        self.logger = logging.getLogger('BrowserManager')
        self._active_instances = {}  # 跟踪活动实例 {task_name: (driver, service)}
        # 添加配置验证日志
        self.logger.info(f"Loaded browser config: {json.dumps(browser_config, indent=2)}")
        '''
        输出浏览器参数
        {
  "driver_path": "F:\\pycode\\chromedriver-win64\\chromedriver.exe",
  "chrome_testing_path": "F:\\pycode\\chrome-win64\\chrome-win64\\chrome.exe",
  "browser_type": "chrome",
  "viewport_width": 1440,
  "quality": 95,
  "headless": true
}
        '''

    @retry(stop=stop_after_attempt(3), wait=wait_fixed(5))
    def create_instance(self, task_name: str) -> (webdriver.Chrome, Service):
        """创建独立浏览器实例（自动重试,增强稳定性版）"""
        options = webdriver.ChromeOptions()
        options.binary_location = self.browser_config['chrome_testing_path']
        options.add_argument(f'--window-size={self.browser_config["viewport_width"]},1080')
        options.add_argument('--disable-gpu')
        options.add_argument('--no-sandbox') # 必须参数
        options.add_argument('--disable-dev-shm-usage') # 共享内存优化
        options.add_argument('--disable-software-rasterizer')
        options.add_argument('--disable-features=VizDisplayCompositor')
        options.add_argument('--force-device-scale-factor=1')
        options.add_argument('--ignore-certificate-errors') # 忽略证书错误

        # 无头模式专用参数
        if self.browser_config['headless']:
            options.add_argument('--headless=new')
            # options.add_argument('--window-size=1440,1080')  # 固定窗口尺寸
            options.add_argument('--disable-extensions')
            options.add_argument('--remote-debugging-pipe')
            # options.add_argument('--remote-debugging-port=9222')  # 固定调试端口

            # 崩溃恢复参数
            options.add_argument('--disable-crash-reporter')
            options.add_argument('--disable-logging')
            options.add_experimental_option('excludeSwitches', ['enable-logging'])

            # 进程隔离配置
            options.add_argument('--disable-setuid-sandbox')
            # options.add_argument('--single-process')  # 单进程模式（慎用可引发浏览器崩溃）

        try:
            service = Service(
                executable_path=self.browser_config['driver_path'],
                service_args=['--verbose', f'--log-path=logs/chromedriver_{task_name}.log'],  # 记录chromedriver日志
                startup_timeout=30  # 延长启动超时
            )

            driver = webdriver.Chrome(service=service,
                                      options=options)
            driver.set_page_load_timeout(60) # 增加页面加载超时
            self._active_instances[task_name] = (driver, service)
            self.logger.info(f"Browser instance created for {task_name}")
            return driver, service
        except Exception as e:
            self.logger.error(f"Create browser failed: {str(e)}")
            # 自动清理可能残留的进程
            if 'service' in locals():
                service.stop()

            # 强制清理残留进程
            os.system("taskkill /f /im chromedriver.exe")
            os.system("taskkill /f /im chrome.exe")
            raise

    def cleanup_instance(self, task_name: str):
        """安全销毁浏览器实例"""
        if task_name in self._active_instances:
            driver, service = self._active_instances.pop(task_name)
            try:
                driver.quit()
                service.stop()
                self.logger.info(f"Browser instance cleaned up for {task_name}")

                """强制清理残留进程"""
                if os.name == 'nt':
                    os.system('taskkill /f /im chromedriver.exe >nul 2>&1')
                    os.system('taskkill /f /im chrome.exe >nul 2>&1')
                else:
                    os.system('pkill -f chromedriver >/dev/null 2>&1')
                    os.system('pkill -f chrome >/dev/null 2>&1')
            except Exception as e:
                self.logger.error(f"Cleanup failed for {task_name}: {str(e)}")

    def cleanup_processes(self):
        """使用WMIC命令精准终止进程"""
        try:
            os.system('wmic process where "name like \'%chromedriver%\'" delete')
            os.system('wmic process where "name like \'%chrome%\'" delete')
        except Exception as e:
            self.logger.warning(f"进程清理失败: {str(e)}")


class CaptureExecutor:
    """原子化截图任务执行器:处理单个截图任务流程"""

    def __init__(self, task_config: Dict[str, Any], browser_config: Dict[str, Any], browser_manager: BrowserManager):
        self.task_config = task_config
        self.browser_config = browser_config  # 存储浏览器配置
        self.browser_manager = browser_manager # 浏览器管理器
        self.logger = logging.getLogger(task_config['name'])

        # 路径标准化处理
        self.temp_dir = Path(task_config['webscroll_capture']['temp_dir'])
        self.save_dir = Path(task_config['webscroll_capture']['save_dir'])
        self.output_path = self.save_dir / task_config['webscroll_capture']['output_path']

        # 确保目录存在
        self.temp_dir.mkdir(parents=True, exist_ok=True)
        self.save_dir.mkdir(parents=True, exist_ok=True)

    def _smart_scroll(self, driver: webdriver.Chrome):
        """医疗报表专用滚动逻辑"""
        last_height = driver.execute_script(
            "return document.getElementById('report1_reportDiv').scrollHeight"
        )
        stable_count = 0

        for _ in range(10):  # 安全循环限制
            driver.execute_script(
                "document.getElementById('report1_reportDiv').scrollTop = 9999999"
            )
            time.sleep(1.5)

            new_height = driver.execute_script(
                "return document.getElementById('report1_reportDiv').scrollHeight"
            )

            if new_height == last_height:
                stable_count += 1
                if stable_count >= 2:
                    break
            else:
                stable_count = 0
                last_height = new_height

            if new_height > 15000:  # 医疗报表高度安全阈值
                break

    def _remove_interfering_elements(self, driver: webdriver.Chrome):
        """移除干扰元素"""
        driver.execute_script("""
            const selectors = [
                'header', 'footer', '.fixed-top', 
                '.cookie-banner', '.ad-container',
                'iframe', '.popup'
            ];

            // 隐藏固定定位元素
            document.querySelectorAll(selectors.join(','))
                .forEach(el => el.style.display = 'none');

            // 禁用动画效果
            const style = document.createElement('style');
            style.textContent = `
                * {
                    animation: none !important;
                    transition: none !important;
                }
            `;
            document.head.appendChild(style);
        """)

    def _force_full_render(self, driver: webdriver.Chrome):
        """强制完整渲染页面所有元素"""
        # 禁用可能影响布局的动态效果
        driver.execute_script("""
            document.body.style.overflow = 'visible';
            document.body.style.transform = 'none';
            const all = document.querySelectorAll('*');
            all.forEach(el => {
                el.style.transition = 'none !important';
                el.style.animation = 'none !important';
            });
        """)
        time.sleep(1)  # 等待样式生效

    def _get_full_dimensions(self, driver: webdriver.Chrome):
        """获取医疗报表容器真实尺寸"""
        return driver.execute_script("""
            const reportDiv = document.getElementById('report1_reportDiv');
            if (!reportDiv) throw new Error('医疗报表容器未找到');

            // 等待最后一行数据渲染完成
            const lastRow = reportDiv.querySelector('tr:last-child');
            lastRow?.scrollIntoView();

            // 计算包含分页信息的精确高度
            const pageMarkers = reportDiv.getElementsByClassName('page-marker');
            const totalHeight = pageMarkers.length > 0 
                ? pageMarkers[pageMarkers.length-1].offsetTop + 120 
                : reportDiv.scrollHeight;

            return [
                reportDiv.scrollWidth + 50,  // 横向留白
                totalHeight
            ];
        """)

    def _set_medical_viewport(self, width, height, driver: webdriver.Chrome):
        """医疗报表视窗安全调整"""
        # 分阶段调整防止内存泄漏
        driver.set_window_size(width + 100, 800)
        time.sleep(0.5)

        # 最终尺寸设置
        driver.set_window_size(width + 100, height + 250)

        # 触发强制重绘
        driver.execute_script("""
            const reportDiv = document.getElementById('report1_reportDiv');
            reportDiv.style.visibility = 'hidden';
            reportDiv.offsetHeight; // 触发重排
            reportDiv.style.visibility = 'visible';
        """)
        time.sleep(1)  # 医疗数据渲染等待

    def _stitch_screenshot(self, output_path: str, driver: webdriver.Chrome):
        """分块截图并拼接"""
        viewport_height = driver.execute_script("return window.innerHeight")
        total_height = driver.execute_script("return document.body.scrollHeight")
        # viewport_height = 1080
        # total_height = 2500
        temp_files = []
        current_position = 0

        try:
            # 分块截图
            while current_position < total_height:
                # 滚动到指定位置
                driver.execute_script(f"window.scrollTo(0, {current_position})")
                time.sleep(0.3)  # 等待渲染

                # 保存临时文件
                temp_file = os.path.join(self.temp_dir, f"part_{current_position}.png")
                driver.save_screenshot(temp_file)
                temp_files.append(temp_file)

                current_position += viewport_height

            # 使用PIL拼接图片
            images = [Image.open(f) for f in temp_files]
            total_width = images[0].width
            combined_image = Image.new(
                'RGB',
                (total_width, total_height),
                color=(255, 255, 255)
            )

            y_offset = 0
            for img in images:
                # 计算实际裁剪区域
                crop_height = min(img.height, total_height - y_offset)
                combined_image.paste(
                    img.crop((0, 0, total_width, crop_height)),
                    (0, y_offset)
                )
                y_offset += crop_height

            # 每天覆盖原来的图片
            save_path = os.path.join(self.save_dir, output_path)
            combined_image.save(save_path, quality=self.browser_config['quality'])
            # print(f"✅ 截图成功保存至: {os.path.abspath(output_path)}")
            print(f"✅ 截图成功保存至: {save_path}")
            self.logger.info(f"[Success] {datetime.now().strftime('%H:%M:%S')} save image:{save_path}")

        except Exception as e:
            raise RuntimeError(f"图片拼接失败: {str(e)}")

    @retry(stop=stop_after_attempt(3),
           wait=wait_exponential(multiplier=1, min=2, max=10),
           retry=retry_if_exception_type((TimeoutException, WebDriverException)),
           before_sleep=lambda retry_state: retry_state.kwargs['self'].browser_manager.cleanup_processes()
           )
    def execute_capture(self) -> bool:
        """执行截图流程（带重试机制）"""
        task_name = self.task_config['name']
        driver = None
        try:
            driver, service = self.browser_manager.create_instance(task_name)

            # 访问目标URL
            driver.get(self.task_config['webscroll_capture']['url'])

            # 等待页面基础加载完成
            WebDriverWait(driver, self.task_config['webscroll_capture']['timeout']).until(
                lambda d: d.execute_script('return document.readyState') == 'complete'
            )

            # 执行智能滚动
            self._smart_scroll(driver)

            # 移除干扰元素
            self._remove_interfering_elements(driver)

            # 强制完整渲染
            self._force_full_render(driver)

            # 获取真实页面尺寸
            width, height = self._get_full_dimensions(driver)
            print(f"页面真实尺寸: {width}x{height}px")

            # 动态调整视口
            self._set_medical_viewport(width, height, driver)

            # 动态调整视口
            # total_height = driver.execute_script("return document.documentElement.scrollHeight")
            # driver.set_window_size(
            #     self.browser_config['viewport_width'],
            #     min(total_height, 16384)  # Chrome最大截图高度限制
            # )

            # 生成临时文件
            timestamp = datetime.now().strftime("%Y%m%d%H%M%S")
            temp_file = self.temp_dir / f"temp_{timestamp}.png"
            driver.save_screenshot(str(temp_file))

            # 格式转换与压缩
            with Image.open(temp_file) as img:
                img.save(self.output_path, quality=self.browser_config['quality'], optimize=True)

            self.logger.info(f"Screenshot saved: {self.output_path}")
            return True

        except Exception as e:
            self.logger.error(f"Capture failed: {str(e)}")
            if driver:
                self.browser_manager.cleanup_instance(task_name)
            return False
        finally:
            if driver:
                self.browser_manager.cleanup_instance(task_name)



class TaskOrchestrator:
    """任务调度协调总控:统一管理所有任务调度"""

    def __init__(self, config: Dict[str, Any]):
        # 添加配置校验
        if not config['global']['browser']:
            raise ValueError("Missing browser config in global section")

        self.config = config
        self.logger = logging.getLogger('TaskOrchestrator')
        self.browser_manager = BrowserManager(config['global']['browser'])
        self.scheduler = BackgroundScheduler()
        self.bot = SyncQyWechatBot(config)
        self.tasks = {}

    def _init_tasks(self):
        """初始化所有任务配置"""
        for task_config in self.config['tasks']:
            task_name = task_config['name']

            # 创建执行器实例
            executor = CaptureExecutor(task_config, self.config['global']['browser'], self.browser_manager) # 传递完整的browser配置节点

            # 注册截图任务
            self._schedule_capture(task_config, executor)

            # 注册发送任务
            self._schedule_send(task_config)

            self.tasks[task_name] = executor

    def _schedule_capture(self, task_config: Dict[str, Any], executor: CaptureExecutor):
        """配置截图触发器"""
        trigger_config = task_config['schedules']['capture']
        job_id = f"{task_config['name']}_capture"

        if trigger_config['type'] == 'cron':
            trigger = CronTrigger(**trigger_config['cron'])
        elif trigger_config['type'] == 'interval':
            trigger = IntervalTrigger(**trigger_config['interval'])

        self.scheduler.add_job(
            executor.execute_capture,
            trigger=trigger,
            id=job_id,
            max_instances=1,
            misfire_grace_time=300,
            coalesce=True,
            next_run_time=datetime.now()  # 启动立即执行一次任务
        )
        self.logger.info(f"Scheduled capture job: {job_id}")

    def _schedule_send(self, task_config: Dict[str, Any]):
        """配置发送触发器"""
        send_config = task_config['schedules']['send']
        job_id = f"{task_config['name']}_send"
        img_path = Path(task_config['webscroll_capture']['save_dir']) / task_config['webscroll_capture']['output_path']

        self.scheduler.add_job(
            self.bot.send_image,
            trigger=CronTrigger(**send_config),
            args=[str(img_path)],
            id=job_id,
            max_instances=1,
            misfire_grace_time=300,
            next_run_time=datetime.now()  # 重启任务，就发送截图
        )
        self.logger.info(f"Scheduled send job: {job_id}")

    def start(self):
        """启动调度系统"""
        # 启动前先清理历史进程
        # self.browser_manager.cleanup_processes()

        self._init_tasks()
        self.scheduler.start()

        # 启动Prometheus监控
        start_http_server(self.config['app']['metrics_port'])

        # 主循环
        try:
            while True:
                time.sleep(3600)  # 降低CPU占用
        except (KeyboardInterrupt, SystemExit):
            self.shutdown()

    def shutdown(self):
        """安全关闭所有服务"""
        self.scheduler.shutdown(wait=False)
        self.bot.shutdown()
        self.browser_manager._active_instances.clear()  # 强制清理残留实例
        self.logger.info("All services terminated")


# --------------------------
# main.py 优化后
# --------------------------
if __name__ == "__main__":
    from secure_config import SecureConfigLoader
    from logger import configure_logger

    # 初始化基础组件
    config = SecureConfigLoader.load("config.yaml")
    configure_logger(config['app']['log_path'])

    # 启动任务协调器
    orchestrator = TaskOrchestrator(config)
    try:
        orchestrator.start()
    except Exception as e:
        logging.critical(f"System crash: {str(e)}")
        orchestrator.shutdown()